'''
- 回测
- python实现简单回测
- 双移动平均策略
- 海龟策略
'''
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns  # seaborn是python中的一个可视化库，是对matplotlib进行二次封装
import akshare as ak

start_date = '2022-01-01'
end_date = '2022-12-31'

stock_002419 = ak.stock_zh_a_hist(symbol="002419",
                                  period="daily",
                                  start_date="20220501",
                                  end_date="20221231", adjust="hfq")

print('------原始数据处理---------')
stock_002419 = pd.DataFrame(stock_002419, columns=['日期', '开盘', '最高', '最低', '收盘', '成交量'])
stock_002419.rename(columns={
    '日期': 'date',
    '开盘': 'open',
    '最高': 'high',
    '最低': 'low',
    '收盘': 'close',
    '成交量': 'volume'
},
    inplace=True)
stock_002419.index = pd.DatetimeIndex(stock_002419['date'])
print(stock_002419.head(10))

'''
1. 创建交易账号
2. 为了不影响原始数据，创建一个新的数据表
'''
# 创建一个新表，只保留原始数据重的日期index
stock_002419_signal = pd.DataFrame(index=stock_002419.index)

# 使用close作为股票价格
stock_002419_signal['price'] = stock_002419['close']

# 增加diff字段，存储股价变化
print('------增加diff字段，存储股价变化，出现NaN---------')
stock_002419_signal['diff'] = stock_002419_signal['price'].diff()
print(stock_002419_signal.head(10))

# 增加diff字段后，第一行会出现空值，使用0来进行填补
print('------NaN用0.0填补------')
stock_002419_signal = stock_002419_signal.fillna(0.0)
print(stock_002419_signal.head(10))

print('------使用signal标记估计上涨或不变：0（卖出），下跌：1(买入)---------')
stock_002419_signal['signal'] = np.where(stock_002419_signal['diff'] >= 0, 0, 1)
print(stock_002419_signal.head(10))

print('------根据交易信号进行下单---------')
print('------卖/买出100股---------')
'''
相对于前一天不变不操作
相对于前一天下跌：买入100股
相对于前一天上涨：卖出100股
'''

stock_002419_signal['order'] = stock_002419_signal['signal'].diff() * 100
stock_002419_signal['hold'] = stock_002419_signal['order'].cumsum()
# stock_002419_signal = stock_002419_signal.fillna('--')
print('------下单情况数据---------')
print(stock_002419_signal.head(10))


'''
回测交易策略
'''
# 初始资金
initial_cash = 20000.00
# 增加trade_stock字段代表 交易股票的市值
stock_002419_signal['trade_stock'] = stock_002419_signal['order'] * stock_002419_signal['price']
# 现金流变化, pd.cumsum()累加函数
# 单次交易金额
stock_002419_signal['use_cash'] = -stock_002419_signal['trade_stock']
# 累加值
stock_002419_signal['use_cash_cum_sum'] = stock_002419_signal['use_cash'].cumsum()
# 初始金额 + 现金流累加（已转换了正负） = 剩余金额；
stock_002419_signal['cash'] = initial_cash + stock_002419_signal['use_cash_cum_sum']
# 持仓市值+现金=总资产
stock_002419_signal['total'] = stock_002419_signal['hold'] * stock_002419_signal['price'] + stock_002419_signal['cash']

print('-------结果--------')
print(stock_002419_signal.head(10))

# 导出表哥
stock_002419_signal.head(30).to_excel('dahua_002419.xls')

'''
用图形表示
总资产
持仓市值
'''
plt.figure(figsize=(10, 6))
plt.plot(stock_002419_signal['total'], '-', label='total assets')  # 总资产
plt.plot(stock_002419_signal['hold'] * stock_002419_signal['price'], '--', label='hold stock')  # 持仓市值

plt.grid()
plt.legend(loc='center right')
plt.show()
